#ifndef _MY_MIXMEMORY_HPP_
#define _MY_MIXMEMORY_HPP_

#include <memory>
#include <cuda_runtime.h>

namespace MMem
{
    class MixMemory
    {
    public:
        MixMemory();
        MixMemory(int device_id);
        MixMemory(const MixMemory& other) = delete;
        MixMemory& operator = (const MixMemory& other) = delete;
        virtual ~MixMemory();

        void set_device(int device_id);
        void set_custream(cudaStream_t* stream);
        void synchronize();

        //传入size，如果size小于等于当前gpu_size_，则直接返回，否则就释放当前gpu资源，重新开辟size大小的资源
        void new_gpu_mem(size_t size);
        //传入size，如果size小于等于当前cpu_size_，则直接返回，否则就释放当前cpu资源，重新开辟size大小的资源
        void new_cpu_mem(size_t size);
        //释放对应设备上申请的gpu内存，指针设为null，size归0
        void release_gpu();
        //释放对应设备上申请的cpu内存，指针设为null，size归0
        void release_cpu();
        //调用release_cpu、release_gpu，先后释放对应设备上申请的cpu、gpu内存，指针设为null，size归0
        void release_all();

        void copy_mem_cpu2cpu(size_t offset_byte, const void* src, size_t num_byte);
        void copy_mem_gpu2gpu(size_t offset_byte, const void* src, size_t num_byte, int device_id);
        void copy_mem_cpu2gpu(size_t offset_byte, const void* src, size_t num_byte, int device_id);
        void copy_mem_gpu2cpu(size_t offset_byte, const void* src, size_t num_byte, int device_id);

        void to_cpu();
        void to_gpu();

        inline int device_id() const{return device_id_;}
        inline bool owner_gpu() const{return gpu_ != nullptr;}//是否拥有gpu资源
        inline bool owner_cpu() const{return cpu_ != nullptr;}//是否拥有cpu资源
        inline size_t gpu_size() const{return gpu_size_;}
        inline size_t cpu_size() const{return cpu_size_;}
        inline void* gpu() const { return gpu_; }
        inline void* cpu() const { return cpu_; }// Pinned Memory
        inline cudaStream_t* custream() {return this->custream_ptr_;}

    private:
        void check_copy_feasible_cpu(size_t offset_byte, size_t num_byte);
        void check_copy_feasible_gpu(size_t offset_byte, size_t num_byte);

    private:
        //cpu指针、cpu数据尺寸
        void* cpu_ = nullptr;
        size_t cpu_size_ = 0;

        //gpu设备号、gpu指针、gpu数据尺寸、gpu流（在类外初始化和销毁）
        int device_id_ = 0;
        void* gpu_ = nullptr;
        size_t gpu_size_ = 0;
        cudaStream_t* custream_ptr_ = nullptr;
    };
}
#endif